# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
# generator_ascend_generate_grouped_matmul_add.py
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul")
class GroupedMatmulAddGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)
        self.m = 0
        self.k = 0
        self.n = 0
        self.e = 0

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        self.m = random.randint(1,2048)
        self.k = random.randint(1,7168)
        self.n = random.randint(1,7168)
        x_shape = case_config.inputs[0][0].shape
        w_shape = case_config.inputs[1][0].shape
        group_list_shape = case_config.inputs[2].shape

        # x shape
        x_shape[0] = self.m # m
        x_shape[1] = self.k # k
        self.e = random.randint(1, min(x_shape[0], 128)) # e

        # weight shape
        w_shape[0] = self.e # e
        w_shape[1] = self.k # k
        w_shape[2] = self.n # n

        #groupList
        group_list_shape[0] = self.e # e
        case_config.inputs[2].range_values = [1, self.m]

        # dtype of x is same as weight's 
        x_dtype = case_config.inputs[0][0].dtype
        case_config.inputs[1][0].dtype = x_dtype

        return case_config
    